#include "common.h"
#include "kernel_operator.h"

namespace AscendC
{
template <typename _DT_DY, typename _DT_X, typename _DT_Z> class Kernel
{
  public:
    TPipe pipe;
    CommonTiling tilingData;
    DefInTensor(DY);
    DefInTensor(X);
    DefOutTensor(Z);

    DefBufVECCALC(TEMP1);
    DefBufVECCALC(TEMP2);

  public:
    __aicore__ inline Kernel()
    {
    }
    __aicore__ inline void Init(GM_ADDR dy, GM_ADDR x, GM_ADDR z, CommonTiling &tiling)
    {
        tilingData = tiling;
        // 切分global
        SetGlobalBuffer(DY, dy, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(X, x, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(Z, z, tilingData.startIdx, tilingData.blockLength);
        // 初始化队列
        InitQueueSimple(DY, tilingData.bufferSize);
        InitQueueSimple(X, tilingData.bufferSize);
        InitQueueSimple(Z, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(TEMP1, tilingData.bufferSize * sizeof(_DT_X));
        InitTBufBuffer(TEMP2, tilingData.bufferSize * sizeof(_DT_X));
    }
    __aicore__ inline void Process()
    {
        auto finnal_progress = tilingData.loopCount - 1;
        for (uint32_t i = 0; i < tilingData.loopCount - 1; ++i)
        {
            CopyIn(i, tilingData.bufferSize);
            Compute(i, tilingData.bufferSize);
            CopyOut(i, tilingData.bufferSize);
        }
        auto copy_count = tilingData.getFinalCopyCount<TypeOf(X)>();
        CopyIn(finnal_progress, copy_count);
        Compute(finnal_progress, tilingData.finalCalcCount);
        CopyOut(finnal_progress, copy_count);
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        EnQueGlobal2Local(DY, index, calcCount);
        EnQueGlobal2Local(X, index, calcCount);
    }
    // x * e^((x - |x|) * 0.851) / (1 + e^(-1.702 * |x|))
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount)
    {
        DeQueSimple(DY);
        DeQueSimple(X);
        QueAllocSimple(Z);
        TBufGet(TEMP1, _DT_X);
        TBufGet(TEMP2, _DT_X);
        // |x|
        Abs(TBufTensorName(TEMP1), LTensorName(X), calcCount);
        
        // -1.702 * |x|
        Muls(TBufTensorName(TEMP2), TBufTensorName(TEMP1), _DT_X(-1.702), calcCount);
        // e^(-1.702 * |x|)
        Exp(LTensorName(Z), TBufTensorName(TEMP2), calcCount);

        // x - |x|
        Sub(TBufTensorName(TEMP2), LTensorName(X), TBufTensorName(TEMP1), calcCount);
        // (x - |x|) * 1.702
        Muls(TBufTensorName(TEMP2), TBufTensorName(TEMP2), _DT_X(1.702), calcCount);
        // e^((x - |x|) * 1.702)
        Exp(TBufTensorName(TEMP1), TBufTensorName(TEMP2), calcCount);
        //  e^(-1.702|x|) +  e^((x - |x|) * 1.702)
        Add(TBufTensorName(TEMP2), TBufTensorName(TEMP1), LTensorName(Z), calcCount);

        // x * e^(-1.702|x|)
        Mul(TBufTensorName(TEMP1), LTensorName(X), LTensorName(Z), calcCount);
        // div_up = e^(-1.702|x|) +  e^((x - |x|) * 1.702) + 1.701 * x * e^(-1.702|x|)
        Axpy(TBufTensorName(TEMP2), TBufTensorName(TEMP1), _DT_X(1.702), calcCount);

        // e^(-1.702 * |x|) + 1
        Adds(TBufTensorName(TEMP1), LTensorName(Z), _DT_X(1), calcCount);
        // div_down = (e^(-1.702 * |x|) + 1) ^ 2
        Mul(LTensorName(Z), TBufTensorName(TEMP1), TBufTensorName(TEMP1), calcCount);

        // res = div_up / div_down
        Div(TBufTensorName(TEMP1), TBufTensorName(TEMP2), LTensorName(Z), calcCount);

        // z = dy * res
        Mul(LTensorName(Z), TBufTensorName(TEMP1), LTensorName(DY), calcCount);
        EnQue(Z);
        QueFree(DY);
        QueFree(X);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        DeQueLocal2Global(Z, index, calcCount);
    }
};
} // namespace AscendC
extern "C" __global__ __aicore__ void fast_gelu_grad(GM_ADDR dy, GM_ADDR x, GM_ADDR z, GM_ADDR workspace,
                                                     GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    CRATE_COMMON_TILING_SIMPLE;
    AscendC::Kernel<DTYPE_DY, DTYPE_X, DTYPE_Z> op;
    op.Init(dy, x, z, commonTiling);
    op.Process();
}